import numpy as np
from stable_baselines3 import PPO
from ymbot_e_pid_env import YmbotEPIDEnv

def target_curve(t):
    curve = np.zeros(30)
    # curve[28] = 20/180*np.pi * np.sin((t-1) * 0.01)  
    curve[29] = 20/180*np.pi * np.sin((t-1) * 0.01)  
    return curve

def main():
    # 创建环境
    env = YmbotEPIDEnv(render_mode='human', window_width=2560, window_height=1440)
    env.set_target_curve(target_curve)  # 设置目标曲线

    # 加载训练好的模型
    model = PPO.load("/home/sh/catkin_ws/src/ymbot_e_control/policy/ppo_ymbot_e_pid.zip")

    # 重置环境
    state, _ = env.reset()
    # print(f"Simulation Timestep: {env.get_simulation_timestep()}")  # 打印仿真步进时间

    # 打开文件以写入数据
    with open('/home/sh/catkin_ws/src/ymbot_e_control/data/ymbot_e_ppo_data.txt', 'w') as file:
        for timestep in range(10000):
            # 使用模型预测动作
            action, _states = model.predict(state, deterministic=True)
            state, reward, done, truncated, info = env.step(action)
            env.render()

            # print(f"Reward: {reward}")  # 打印奖励
            # print(f"timestep: {env.timestep}")  
            # print(f"action[24:30]: {action[24:30]}")
            # print(f"action[18:24: {action[18:24]}")
            print(f"action: {action}")
            
            # 导出倒数六个关节的数据
            data = [timestep]
            current_positions = []
            target_positions = []
            kps = []
            kis = []
            kds = []
            torques = []
            for i in range(24, 30):
                current_positions.append(state[i])
                target_positions.append(env.setpoint_curve(env.timestep)[i])
                kps.append(0)
                kis.append(0)
                kds.append(0)
                torques.append(env.control_output[i])
            data.extend(current_positions)
            data.extend(target_positions)
            data.extend(kps)
            data.extend(kis)
            data.extend(kds)
            data.extend(torques)
            file.write(','.join(map(str, data)) + '\n')
            
            # if done:
            #     state, _ = env.reset()  # 重新初始化环境状态

    env.close()

if __name__ == "__main__":
    main()
